import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import LSTM, Dense, Embedding
from tensorflow.keras.preprocessing.text import Tokenizer, tokenizer_from_json
from tensorflow.keras.preprocessing.sequence import pad_sequences
import re
import json
import os
from datetime import datetime

class CodeGeneratorAI:
    def __init__(self, max_sequence_length=200, vocab_size=5000):
        self.max_sequence_length = max_sequence_length
        self.vocab_size = vocab_size
        self.tokenizer = Tokenizer(num_words=vocab_size, filters='')
        self.model = None

    def preprocess_code(self, code_samples):
        processed = []
        for code in code_samples:
            code = re.sub(r'#.*', '', code)  # 移除注释
            processed.append(f'<START> {code} <END>')
        return processed

    def train(self, code_samples, epochs=150, batch_size=32):
        processed_code = self.preprocess_code(code_samples)
        self.tokenizer.fit_on_texts(processed_code)
        total_words = len(self.tokenizer.word_index) + 1
        
        # 创建训练序列
        input_sequences = []
        for line in processed_code:
            token_list = self.tokenizer.texts_to_sequences([line])[0]
            for i in range(1, len(token_list)):
                n_gram_sequence = token_list[:i+1]
                input_sequences.append(n_gram_sequence)
        
        input_sequences = pad_sequences(input_sequences, maxlen=self.max_sequence_length, padding='pre')
        predictors, label = input_sequences[:,:-1], input_sequences[:,-1]
        label = tf.keras.utils.to_categorical(label, num_classes=total_words)
        
        # 定义模型
        self.model = Sequential()
        self.model.add(Embedding(total_words, 100, input_length=self.max_sequence_length-1))
        self.model.add(LSTM(150, return_sequences=True))
        self.model.add(LSTM(100))
        self.model.add(Dense(total_words, activation='softmax'))
        
        self.model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
        self.model.fit(predictors, label, epochs=epochs, batch_size=batch_size, verbose=1)

    def generate_code(self, prompt, max_length=200, temperature=0.6):
        if self.model is None:
            raise Exception("模型尚未训练，请先调用train方法")
            
        generated = f'<START> {prompt}'
        for _ in range(max_length):
            token_list = self.tokenizer.texts_to_sequences([generated])[0]
            token_list = pad_sequences([token_list], maxlen=self.max_sequence_length-1, padding='pre')
            
            predictions = self.model.predict(token_list, verbose=0)[0]
            predictions = np.asarray(predictions).astype('float64')
            predictions = np.log(predictions) / temperature
            exp_preds = np.exp(predictions)
            predictions = exp_preds / np.sum(exp_preds)
            probas = np.random.multinomial(1, predictions, 1)
            predicted_index = np.argmax(probas)
            
            output_word = ""
            for word, index in self.tokenizer.word_index.items():
                if index == predicted_index:
                    output_word = word
                    break
                    
            if output_word == '<END>':
                break
                
            generated += " " + output_word
            
        generated = generated.replace('<START> ', '')
        generated = generated.replace(' <END>', '')
        return generated

    def save_model(self, model_path='code_model.h5', tokenizer_path='tokenizer_config.json'):
        self.model.save(model_path)
        with open(tokenizer_path, 'w') as f:
            f.write(self.tokenizer.to_json())

    def load_model(self, model_path='code_model.keras', tokenizer_path='tokenizer_config.json'):
        self.model = load_model(model_path)
        with open(tokenizer_path, 'r') as f:
            self.tokenizer = tokenizer_from_json(f.read())


class InteractiveCodeEditor:
    def __init__(self, code_generator):
        self.code_generator = code_generator
        self.current_code = ""
        self.history = []

    def generate_from_prompt(self, prompt, temperature=0.6):
        """根据提示生成代码"""
        print(f"正在根据提示生成代码...")
        self.current_code = self.code_generator.generate_code(prompt, temperature=temperature)
        self.history.append(("生成", prompt, self.current_code))
        return self.current_code

    def refine_code(self, instruction, temperature=0.6):
        """基于当前代码和指令进行优化"""
        if not self.current_code:
            print("没有当前代码，请先生成代码")
            return
        
        print(f"正在根据指令优化代码...")
        # 构建提示：当前代码 + 指令
        prompt = f"{self.current_code}\n\n# 根据以下指令优化代码:\n# {instruction}\n"
        new_code = self.code_generator.generate_code(prompt, temperature=temperature)
        
        # 智能提取优化部分（假设模型在原有代码基础上继续生成）
        if new_code.startswith(self.current_code):
            refined_part = new_code[len(self.current_code):].strip()
            self.current_code = f"{self.current_code}\n{refined_part}"
        else:
            self.current_code = new_code
            
        self.history.append(("优化", instruction, self.current_code))
        return self.current_code

    def edit_code(self, new_code):
        """手动编辑代码"""
        self.current_code = new_code
        self.history.append(("手动编辑", "用户修改", self.current_code))
        return self.current_code

    def save_code(self, filename=None):
        """保存当前代码到文件"""
        if not filename:
            filename = f"generated_code_{datetime.now().strftime('%Y%m%d_%H%M%S')}.py"
            
        with open(filename, 'w') as f:
            f.write(self.current_code)
            
        print(f"代码已保存到: {filename}")
        return filename

    def show_history(self):
        """显示操作历史"""
        for i, (action, instruction, _) in enumerate(self.history):
            print(f"{i+1}. [{action}] {instruction[:50]}..." if len(instruction) > 50 else instruction)


def main():
    # 初始化模型
    print("正在加载代码生成模型...")
    ai = CodeGeneratorAI()
    
    # 检查是否有预训练模型
    if os.path.exists('code_model.h5') and os.path.exists('tokenizer_config.json'):
        ai.load_model()
        print("已加载预训练模型")
    else:
        print("未找到预训练模型，请先训练模型（运行train_model.py）")
        return
    
    # 初始化编辑器
    editor = InteractiveCodeEditor(ai)
    
    print("\n=== 交互式代码生成器 ===")
    print("输入 'g' + 提示词: 根据提示生成代码")
    print("输入 'r' + 指令: 基于当前代码优化")
    print("输入 'e': 手动编辑代码")
    print("输入 's': 保存代码")
    print("输入 'h': 查看历史")
    print("输入 'q': 退出")
    
    while True:
        command = input("\n> ").strip()
        
        if command.lower() == 'q':
            break
            
        elif command.lower().startswith('g'):
            prompt = command[1:].strip()
            if not prompt:
                print("请提供生成提示词")
                continue
                
            code = editor.generate_from_prompt(prompt)
            print("\n生成的代码:")
            print(code)
            
        elif command.lower().startswith('r'):
            instruction = command[1:].strip()
            if not instruction:
                print("请提供优化指令")
                continue
                
            code = editor.refine_code(instruction)
            print("\n优化后的代码:")
            print(code)
            
        elif command.lower() == 'e':
            print("\n当前代码:")
            print(editor.current_code)
            print("\n请输入新代码（输入'done'结束编辑）:")
            
            lines = []
            while True:
                line = input()
                if line.lower() == 'done':
                    break
                lines.append(line)
                
            new_code = '\n'.join(lines)
            editor.edit_code(new_code)
            print("\n更新后的代码:")
            print(editor.current_code)
            
        elif command.lower() == 's':
            filename = input("请输入文件名（留空自动生成）: ").strip()
            editor.save_code(filename if filename else None)
            
        elif command.lower() == 'h':
            print("\n操作历史:")
            editor.show_history()
            
        else:
            print("未知命令，请重新输入")


if __name__ == "__main__":
    main()